import { ListenableGenerator } from "./ListenableGenerator"

export class GeneratorReuseManager {
	currentGenerator: ListenableGenerator<string> | undefined
	pendingGeneratorPrefix: string | undefined
	pendingCompletion = ""

	constructor(private readonly onError: (err: unknown) => void) {}

	private _createListenableGenerator(abortController: AbortController, gen: AsyncGenerator<string>, prefix: string) {
		this.currentGenerator?.cancel()

		const listenableGen = new ListenableGenerator(gen, this.onError, abortController)
		listenableGen.listen((chunk) => (this.pendingCompletion += chunk ?? ""))

		this.pendingGeneratorPrefix = prefix
		this.pendingCompletion = ""
		this.currentGenerator = listenableGen
	}

	private shouldReuseExistingGenerator(prefix: string): boolean {
		return (
			!!this.currentGenerator &&
			!!this.pendingGeneratorPrefix &&
			(this.pendingGeneratorPrefix + this.pendingCompletion).startsWith(prefix) &&
			// for e.g. backspace
			this.pendingGeneratorPrefix?.length <= prefix?.length
		)
	}

	async *getGenerator(
		prefix: string,
		newGenerator: (abortSignal: AbortSignal) => AsyncGenerator<string>,
		multiline: boolean,
	): AsyncGenerator<string> {
		// If we can't reuse, then create a new generator
		if (!this.shouldReuseExistingGenerator(prefix)) {
			// Create a wrapper over the current generator to fix the prompt
			const abortController = new AbortController()
			this._createListenableGenerator(abortController, newGenerator(abortController.signal), prefix)
		}

		// Already typed characters are those that are new in the prefix from the old generator
		let typedSinceLastGenerator = prefix.slice(this.pendingGeneratorPrefix?.length) || ""
		for await (let chunk of this.currentGenerator?.tee() ?? []) {
			if (!chunk) {
				continue
			}

			// Ignore already typed characters in the completion
			while (chunk.length && typedSinceLastGenerator.length) {
				if (chunk[0] === typedSinceLastGenerator[0]) {
					typedSinceLastGenerator = typedSinceLastGenerator.slice(1)
					chunk = chunk.slice(1)
				} else {
					break
				}
			}

			// Break at newline unless we are in multiline mode
			const newLineIndex = chunk.indexOf("\n")
			if (newLineIndex >= 0 && !multiline) {
				yield chunk.slice(0, newLineIndex)
				break
			} else if (chunk !== "") {
				yield chunk
			}
		}
	}
}
